package org.finesys.common.log.trace.interceptor.gateway;

import java.util.HashMap;
import java.util.Map;
import java.util.Optional;

import org.reactivestreams.Subscription;
import org.slf4j.MDC;
import org.springframework.cglib.beans.BeanMap;

import reactor.core.CoreSubscriber;
import reactor.util.context.Context;

public class MDCSubscriber<T> implements CoreSubscriber<T> {

    private final static String TRACE_ID = "traceId";
    private static final String SKYWALKING_CTX_SNAPSHOT = "SKYWALKING_CONTEXT_SNAPSHOT";
    private final CoreSubscriber<Object> actual;

    public MDCSubscriber(CoreSubscriber<Object> actual) {
        this.actual = actual;
    }

    @Override
    public void onSubscribe(Subscription s) {
        actual.onSubscribe(s);
    }

    @Override
    public void onNext(Object o) {
        Context c = actual.currentContext();

        Optional<String> traceIdOptional = Optional.empty();
        if (!c.isEmpty() && c.hasKey(SKYWALKING_CTX_SNAPSHOT)) {
            traceIdOptional = Optional.of(c.get(SKYWALKING_CTX_SNAPSHOT)).map(MDCSubscriber::beanToMap).map(t -> t.get(TRACE_ID)).map(Object::toString);
        }
        try (MDC.MDCCloseable mdcCloseable = MDC.putCloseable(TRACE_ID, traceIdOptional.orElse("N/A"))) {
            actual.onNext(o);
        }
    }

    @Override
    public void onError(Throwable t) {
        actual.onError(t);
    }

    @Override
    public void onComplete() {
        actual.onComplete();
    }

    @Override
    public Context currentContext() {
        return actual.currentContext();
    }

    public static <T> Map<String, Object> beanToMap(T bean) {
        Map<String, Object> map = new HashMap<>();
        if (bean != null) {
            BeanMap beanMap = BeanMap.create(bean);
            for (Object key : beanMap.keySet()) {
                if (TRACE_ID.equals(key)) {
                    map.put(TRACE_ID, beanMap.get(TRACE_ID));
                }
            }
        }
        return map;
    }
}
